import h5py
import numpy as np
from typing import Dict, Optional, Any

class DataKeys:
    """HDF5文件中的数据键映射"""
    MAPPINGS = {
        "DTL-010:RFS-DIG-101:Dwn0-XAxis": "axis",
        "MEBT-010:RFS-DIG-101:Dwn0-Cmp0": "mag_cav",
        "MEBT-010:RFS-DIG-101:Dwn0-Cmp1": "phase_cav",
        "MEBT-010:RFS-DIG-101:Dwn5-Cmp0": "mag_for",
        "MEBT-010:RFS-DIG-101:Dwn5-Cmp1": "phase_for",
        "MEBT-010:RFS-DIG-101:Dwn6-Cmp0": "mag_refl",
        "MEBT-010:RFS-DIG-101:Dwn6-Cmp1": "phase_refl"
    }

class Config:
    """数据处理配置"""
    EFFECTIVE_LENGTH = 0.855  # 有效长度
    POWER_SCALE = 1000       # 功率缩放因子
    VOLTAGE_SCALE = 1000     # 电压缩放因子

def extract_dataset(group: h5py.Group, index: int) -> Optional[Any]:
    """从HDF5组中提取数据集
    
    Args:
        group: HDF5数据组
        index: 时间戳索引
        
    Returns:
        Optional[Any]: 提取的数据集，如果不存在则返回None
    """
    if len(group) == 0:
        return None
        
    if index >= len(group):
        print(f"警告：索引 {index} 超出范围，组中只有 {len(group)} 个数据集")
        return None
        
    time_stamps = list(group)
    return group[time_stamps[index]][()]

def validate_raw_powers(raw_powers: Dict) -> None:
    """验证原始数据的完整性"""
    required_keys = [
        "mag_cav", "phase_cav",
        "mag_for", "phase_for",
        "mag_refl", "phase_refl"
    ]
    
    missing_keys = [key for key in required_keys if key not in raw_powers]
    if missing_keys:
        raise KeyError(f"缺少必要的键: {', '.join(missing_keys)}")

def calibrate_magnitudes(mag_cav, mag_for, mag_refl, RQ, Qext, beta):
    """校准幅度值"""
    # 腔体幅度校准
    cav_cal = mag_cav
    
    # 前向和反射功率校准
    power_factor = RQ * Qext * (1 + beta) / beta
    scale = Config.POWER_SCALE
    
    for_cal = np.sqrt(np.abs(mag_for * scale * power_factor)) / Config.VOLTAGE_SCALE
    refl_cal = np.sqrt(np.abs(mag_refl * scale * power_factor)) / Config.VOLTAGE_SCALE
    
    return cav_cal, for_cal, refl_cal

def convert_to_iq(magnitude, phase):
    """将幅度和相位转换为IQ分量"""
    phase_rad = np.deg2rad(phase)
    return (
        magnitude * np.cos(phase_rad),
        magnitude * np.sin(phase_rad)
    )

def read_and_process_hdf5(hdf5_filename: str, index: int, RQ: float, Qext: float, beta: float) -> Dict[str, Any]:
    """读取HDF5文件并处理数据
    
    Args:
        hdf5_filename: HDF5文件路径
        index: 时间戳索引
        RQ: 电阻品质因数乘积
        Qext: 外部品质因数
        beta: 耦合系数
        
    Returns:
        Dict[str, Any]: 包含处理后的IQ数据的字典
    """
    try:
        # 1. 读取原始数据
        raw_powers = {}
        with h5py.File(hdf5_filename, 'r') as hdf5_file:
            for h5_key, raw_key in DataKeys.MAPPINGS.items():
                if h5_key not in hdf5_file:
                    print(f"警告：文件中缺少键 {h5_key}")
                    continue
                    
                group = hdf5_file[h5_key]
                if raw_key == "axis":
                    raw_powers[raw_key] = extract_dataset(group, 0)
                else:
                    raw_powers[raw_key] = extract_dataset(group, index)

        # 2. 验证数据完整性
        validate_raw_powers(raw_powers)
        
        # 3. 校准幅度
        mag_cav_cal, mag_for_cal, mag_refl_cal = calibrate_magnitudes(
            raw_powers["mag_cav"],
            raw_powers["mag_for"],
            raw_powers["mag_refl"],
            RQ, Qext, beta
        )
        
        # 4. 转换为IQ分量
        Vcav_I, Vcav_Q = convert_to_iq(mag_cav_cal, raw_powers["phase_cav"])
        Vfor_I, Vfor_Q = convert_to_iq(mag_for_cal, raw_powers["phase_for"])
        Vrefl_I, Vrefl_Q = convert_to_iq(mag_refl_cal, raw_powers["phase_refl"])
        
        # 5. 返回处理后的数据
        return {
            "Vcav_I": Vcav_I, "Vcav_Q": Vcav_Q,
            "Vfor_I": Vfor_I, "Vfor_Q": Vfor_Q,
            "Vrefl_I": Vrefl_I, "Vrefl_Q": Vrefl_Q
        }
        
    except FileNotFoundError:
        raise FileNotFoundError(f"找不到文件：{hdf5_filename}")
    except Exception as e:
        raise ValueError(f"处理数据时出错：{str(e)}")

if __name__ == '__main__':
    # 测试代码
    try:
        test_file = "path/to/test.hdf5"
        data = read_and_process_hdf5(test_file, 1, RQ=8893, Qext=76.7, beta=1.2)
        print("成功处理数据：", list(data.keys()))
    except Exception as e:
        print(f"测试失败：{str(e)}")        